import pandas as pd
import numpy as np
import torch
from torch import nn

train_features = np.load('./data/train_features.npy', allow_pickle=True)
train_labels = np.load('./data/train_labels.npy', allow_pickle=True)
train_features = torch.tensor(train_features.astype(np.float32), dtype=torch.float)
train_labels = torch.tensor(train_labels)

num_inputs = 21
num_outputs = 39


num_epochs = 5
batch_size = 32
lr = 0.1

dataset = torch.utils.data.TensorDataset(train_features, train_labels)
train_iter = torch.utils.data.DataLoader(dataset, batch_size, shuffle=True)

net = nn.Sequential()
net.add_module('Dense',nn.Linear(num_inputs, num_outputs))
# torch框架会自动初始化模型参数
#可以用如下代码查看参数
#for parame in net.named_parameters():
#    print(parame)

loss = nn.CrossEntropyLoss()
# torch框架中CrossEntropyLoss自带Softmax运算，所以网络部分没有Softmax层
optimizer = torch.optim.SGD(net.parameters(), lr=lr)

for epoch in range(num_epochs):
    train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
    for X, y in train_iter:
        y_hat = net(X)
        l = loss(y_hat, y).sum()
        optimizer.zero_grad()
        l.backward()
        optimizer.step()
        train_l_sum += l.item()
        train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
        n += y.shape[0]
    print('epoch %d, loss %.4f, train acc %.3f'
          % (epoch + 1, train_l_sum / n, train_acc_sum / n))

net.add_module('Softmax', nn.Softmax(dim=-1))
test_features = np.load('./data/test_features.npy', allow_pickle=True)
test_features = torch.tensor(test_features.astype(np.float32), dtype=torch.float)
testResult = net(test_features).detach().numpy()
sampleSubmission = pd.read_csv('input/sf-crime/sampleSubmission.csv.zip')
Result_pd = pd.DataFrame(testResult,
                         index=sampleSubmission.index,
                         columns=sampleSubmission.columns[1:])
Result_pd.to_csv('working/sampleSubmission(v0.1).csv', index_label='Id')



if __name__ == '__main__':
    print()